#include <chrono>
#include <iomanip>
#include <iostream>
#include <ostream>
#include <vector>

#include <cstring>
#include <fstream>
#include <numeric>

#include "paddle/pdmodel_manager.h"
#include "base_util/utils.h"


int main(int argc, char **argv) {

  // det /det/car **/model **.jpg
  std::string rec_type = argv[1];
  std::string model_dir = argv[2];
  std::string img_path = argv[3];

  LogInfo *log_info = new LogInfo(7, "./log");
  ai::PDModelManager *pdmodel_manager = new ai::PDModelManager(log_info);
  std::cout << "new PDModelManager ok..." << std::endl;

  int ret = 0;
	std::vector<ai::InitModelData> init_datas; 
	// init_datas.push_back(InitModelData(model_path+"/det_egg", model_det, 0));
	// init_datas.push_back(InitModelData(model_path+"/det_ringelman", model_det, 0));
	// init_datas.push_back(InitModelData(model_path+"/det_handwriting", model_det, 0));

	// init_datas.push_back(InitModelData(model_path+"/det", model_det, 0));
	// init_datas.push_back(InitModelData(model_path+"/rec", 1));

  // init_datas.push_back(InitModelData(model_path+"/ocr/det", model_text_db, 0));
  // init_datas.push_back(InitModelData(model_path+"/ocr/rec", model_text_rec, 1));
  // // init_datas.push_back(InitModelData(model_path+"/ocr/cls", model_cls, 2));

  // init_datas.push_back(InitModelData(model_path+"/cls/car", model_cls, 0));
	// init_datas.push_back(InitModelData(model_path+"/seg/sand",     model_seg, 0));
  if (rec_type == "cls") {
	  init_datas.push_back(ai::InitModelData(model_dir, "./cls/weather", "weather", ai::model_cls, 0));
  }
  else if (rec_type == "det") {
	  init_datas.push_back(ai::InitModelData(model_dir, "./det/car", "car", ai::model_det, 0));
  }
  else if (rec_type == "seg") {
	  init_datas.push_back(ai::InitModelData(model_dir, "./seg/road", "road", ai::model_seg, 0));
  }
	// init_datas.push_back(InitModelData(model_path+"/cls/weather", model_cls, 0));

	if ((ret = pdmodel_manager->init_model(init_datas))!=0){
    std::cout << "init_model: " << ret << std::endl;
    return ret;
  }
  std::cout << "init_model ok..." << std::endl;

  auto start = std::chrono::system_clock::now();


  std::vector<bf::path> files;
  if (bf::is_directory(img_path)) {
    ret = get_files_in_dir(img_path, files, ".jpg");
  } else {
    files.push_back(bf::path(img_path));
  }
  std::cout << "files: " << files.size() << std::endl;
  // // cls batch=1
  if (rec_type == "cls") {
    for (auto f : files) {
      cv::Mat img;
      img = cv::imread(f.string(), cv::IMREAD_COLOR);
      std::vector<ai::BaseInfo*> cls_info;
      int64 t = cv::getTickCount();
      ret = pdmodel_manager->cls(img, cls_info, 0);
      int64 end = cv::getTickCount();
      float ms = static_cast<float>(end - t) / cv::getTickFrequency();
      std::cout <<  "res: " << ((ai::ClsInfo*)(cls_info[0]))->category <<" time: " << ms <<" " << f.string() << std::endl;
    }

  }
  // // det
  if (rec_type == "det") {
    for (auto f : files) {
      // std::string f = argv[2];
      int64 t = cv::getTickCount();
      cv::Mat srcimg = cv::imread(f.string(), cv::IMREAD_COLOR);
      std::vector<ai::BaseInfo*> det_info;
      ret = pdmodel_manager->det(srcimg, det_info, 0);
      int64 end = cv::getTickCount();
      float ms = static_cast<float>(end - t) / cv::getTickFrequency();
      std::cout <<  "det time: " << ms <<" " << f.string() << std::endl;

      for (auto box : det_info) {
        ai::DetInfo* dP = (ai::DetInfo*)box;
        cv::rectangle(srcimg, dP->det_box, cv::Scalar(255, 0, 255), 2);
      }
      cv::imwrite(f.string()+ ".draw.jpg", srcimg);

      std::cout << "ret: " << ret << std::endl;
      if (ret!=0){return ret;}
    }
  }

  // // text rec
  if (rec_type == "text_rec") {
    int ret = 0;
    // det
    std::string img_path(argv[2]);
    cv::Mat cv_image = cv::imread(img_path); 
    // rec
    std::vector<ai::BaseInfo*> rec_ifo;
    ret = pdmodel_manager->text_rec_lite(cv_image, rec_ifo, 1);
    if (ret != 0){return ret;}
    std::cout << rec_ifo.size() << std::endl;
    std::cout << "text: " << ((ai::TextRecInfo*)(rec_ifo[0]))->text << std::endl;
    std::cout << "score: " <<((ai::TextRecInfo*)(rec_ifo[0]))->score << std::endl;
  }

  // // cls
  // {
  //   std::vector<cv::Mat> imgs;
  //   imgs.push_back(cv::imread(argv[2], cv::IMREAD_COLOR));
  //   imgs.push_back(cv::imread(argv[3], cv::IMREAD_COLOR));
  //   imgs.push_back(cv::imread(argv[4], cv::IMREAD_COLOR));
  //   imgs.push_back(cv::imread(argv[5], cv::IMREAD_COLOR));
  //   std::vector<ai::ClsInfo> cls_infos;
  //   std::cout << "before obj_det..." << std::endl;
  //   ret = pdmodel_manager->cls(imgs, cls_infos, 0);
  //   std::cout << "cls_infos.size(): " << cls_infos.size() << std::endl;
  //   std::cout << "ret: " << ret << std::endl;
  // }

  // // seg
  // {  
  //   std::string img_path(argv[2]);
  //   cv::Mat srcimg = cv::imread(img_path, cv::IMREAD_COLOR);
  //   std::cout << srcimg.cols << " " << srcimg.rows << std::endl;
  //   std::vector<ai::DetInfo> text_lines;
  //   std::cout << "before seg_det..." << std::endl;
  //   ret = pdmodel_manager->seg_det(&srcimg, text_lines, 0);
  //   std::cout << "ret: " << ret << std::endl;
  //   if (ret!=0){return ret;}
  // }

  // // seg
  // {
  //   std::string img_path(argv[2]);
  //   cv::Mat srcimg = cv::imread(img_path, cv::IMREAD_COLOR);


  //   std::vector<int> xs = {138,150,163,174,187,208,252,296,316,308,286,263,242,219,163,156,138};
  //   std::vector<int> ys = {194,214,231,249,266,267,260,249,246,230,210,193,183,172,183,186,194};
  //   std::vector<cv::Point> points;
  //   for (int i = 0; i < xs.size(); i++) { points.push_back(cv::Point(xs[i], ys[i])); }
  //   auto minmax_x = std::minmax_element(xs.begin(), xs.end());
  //   auto minmax_y = std::minmax_element(ys.begin(), ys.end());
  //   cv::Rect roi = cv::Rect(cv::Point(*(minmax_x.first), *(minmax_y.first)), 
  //                           cv::Point(*(minmax_x.second), *(minmax_y.second)));
  //   std::cout << "roi: " << roi.x << " " << roi.y << " " << roi.width << " " << roi.height << std::endl;
  //   cv::Mat fg = cv::Mat::zeros(srcimg.size(), CV_8UC3);
  //   cv::Mat mask = cv::Mat::zeros(srcimg.size(), CV_8UC1);
  //   cv::fillConvexPoly(mask, points, cv::Scalar(255)); 
  //   cv::imwrite("mask.jpg", mask);

  //   cv::copyTo(srcimg, fg, mask);
  //   cv::imwrite("fg.jpg", fg);

  //   cv::Mat img = fg(roi).clone();
  //   cv::imwrite("img.jpg", img);


  //   std::vector<ai::DetInfo> text_lines;
  //   std::cout << "before seg_det..." << std::endl;
  //   ret = pdmodel_manager->seg_det(&img, text_lines, 0);
  //   std::cout << "ret: " << ret << std::endl;
  //   if (ret!=0){return ret;}

  //   float weight = 0.6;
  //   cv::Mat color_map =  cv::Mat::zeros(srcimg.size(), CV_8UC3) + cv::Scalar(0,255,0);
  //   cv::imwrite("color_map.jpg", color_map);
  //   for (int i = 0; i < text_lines.size(); i++) {
  //     if (text_lines[i].class_idx != 2) { continue; }
  //     cv::Mat cur_mask = cv::Mat::zeros(srcimg.size(), CV_8UC1);
  //     cv::imwrite(std::to_string(i) + "_mask.jpg", text_lines[i].mask);

  //     text_lines[i].mask.copyTo(cur_mask(roi));
  //     cv::imwrite(std::to_string(i) + "_cur_mask.jpg", cur_mask);

  //     cv::Mat merged;
  //     cv::addWeighted(srcimg, weight, color_map, 1-weight, 0, merged);
  //     cv::imwrite(std::to_string(i) + "_merged.jpg", merged);

  //     cv::copyTo(merged, srcimg, cur_mask);
  //     cv::imwrite(std::to_string(i) + "_srcimg.jpg", srcimg);
  //   }
  //   std::vector<std::vector<cv::Point> > pointss = {points};
  //   cv::drawContours(srcimg, pointss, 0, cv::Scalar(0,255,255), 1);
  //   cv::imwrite("srcimg.jpg", srcimg);
  // }

  

  // // det batch
  // {
  //   for (int i = 0; i < 3; i++) {
  //     std::vector<cv::Mat> imgs;
  //     imgs.push_back(cv::imread(argv[2], cv::IMREAD_COLOR));
  //     imgs.push_back(cv::imread(argv[3], cv::IMREAD_COLOR));
  //     imgs.push_back(cv::imread(argv[4], cv::IMREAD_COLOR));
  //     imgs.push_back(cv::imread(argv[5], cv::IMREAD_COLOR));
  //     std::vector<std::vector<ai::DetInfo>> det_infos;
  //     std::cout << "before obj_det..." << std::endl;
  //     ret = pdmodel_manager->obj_det(imgs, det_infos, 0);
  //     std::cout << "ret: " << ret << std::endl;
  //     {
  //       std::cout << "imgs.size(): " << imgs.size() << std::endl;
  //       std::cout << "det_infos.size(): " << det_infos.size() << std::endl;
  //       for (int i = 0; i < imgs.size(); i++) {
  //         for (auto box : det_infos[i]) {
  //           cv::rectangle(imgs[i], box.det_box, cv::Scalar(255, 0, 255), 2);
  //         }
  //         cv::imwrite(std::to_string(i) + "_.jpg", imgs[i]);
  //       }
  //     }
  //     if (ret!=0){return ret;}
  //   }
  // }

  // egg pic
  // {
  //   std::string video_path(argv[2]);
  //   std::string prefix_path = bf::path(video_path).parent_path().append( \
  //     bf::path(video_path).stem().string()).string();
  //   // if (!bf::exists(prefix_path)){
  //   //   bf::create_directories(prefix_path);
  //   // }
  //   cv::Mat img = cv::imread(video_path);
  //   std::vector<ai::DetInfo> det_info;
  //   ret = pdmodel_manager->obj_det(img, det_info, 0);
  //   for (auto &b : det_info) {
  //     // cv::imwrite(prefix_path +"/" +std::to_string(count++) +".jpg", 
  //     //   img(b.det_box));
  //     cv::Mat crop = img(b.det_box);
  //     ClasInfoPD cla_info;
  //     ret = pdmodel_manager->cls(crop, cla_info, 1);
  //     cv::rectangle(img, b.det_box, 
  //       cla_info.class_idx == 0 ? cv::Scalar(0, 0, 255) : cv::Scalar(0, 255, 0), 
  //       2);
  //     //  cla_info.class_idx == 0? "NG":"OK"  std::to_string(b.det_score)
  //     cv::putText(img, cla_info.class_idx == 0? "NG":"OK", \
  //       cv::Point(b.det_box.x, b.det_box.y - 17), cv::FONT_HERSHEY_SIMPLEX, \
  //       1, 
  //       cla_info.class_idx == 0 ? cv::Scalar(0, 0, 255) : cv::Scalar(0, 255, 0) ,
  //       2);
  //   }
  //   cv::imwrite(prefix_path + "_rec.jpg", img);
  // }

  // // ringelman pic
  // {
  //   std::string img_path(argv[2]);
  //   std::string prefix_path = bf::path(img_path).parent_path().append( \
  //     bf::path(img_path).stem().string()).string();
  //   cv::Mat img = cv::imread(img_path);
  //   std::vector<ai::DetInfo> det_info;
  //   ret = pdmodel_manager->obj_det(&img, det_info, 0);
  //   for (auto &b : det_info) {
  //     // cv::imwrite(prefix_path +"/" +std::to_string(count++) +".jpg", 
  //     //   img(b.det_box));
  //     cv::rectangle(img, b.det_box, cv::Scalar(0, 255, 0), 2);
  //   }
  //   cv::imwrite(prefix_path + "_rec.jpg", img);
  // }

  // // egg det video
  // {
  //   std::string video_path(argv[2]);
  //   std::string prefix_path = bf::path(video_path).parent_path().append( \
  //     bf::path(video_path).stem().string()).string();
  //   // if (!bf::exists(prefix_path)){
  //   //   bf::create_directories(prefix_path);
  //   // }
  //   cv::VideoCapture cap;
  //   if (!cap.open(video_path)){
  //     return 100;
  //   }
  // 	int width = cap.get(3);
	//   int height = cap.get(4);
  //   cv::VideoWriter writer;
  //   bool is_opened = writer.open(prefix_path + "_output.avi", cv::VideoWriter::fourcc('M','J','P','G'),
  //   20, cv::Size(width, height), true );
  //   if (!writer.isOpened()) {return 100;}
  //   int count = 0;
  //   bool ok = true;
  //   while (ok) {
  //     cv::Mat img;
  //     cap >> img;
  //     if (img.empty()){
  //       ok = false;
  //       break;
  //     }
  //     count++;
  //     std::vector<ai::DetInfo> det_info;
  //     ret = pdmodel_manager->obj_det(img, det_info, 0);
  //     for (auto &b : det_info) {
  //       // cv::imwrite(prefix_path +"/" +std::to_string(count++) +".jpg", 
  //       //   img(b.det_box));
  //       cv::Mat crop = img(b.det_box);
  //       ClasInfoPD cla_info;
  //       std::cout << count << "  " << crop.cols << "  " << crop.rows << std::endl;
  //       ret = pdmodel_manager->cls(crop, cla_info, 1);
  //       cv::rectangle(img, b.det_box, 
  //         cla_info.class_idx == 0 ? cv::Scalar(0, 0, 255) : cv::Scalar(0, 255, 0), 
  //         2);
  //       //  cla_info.class_idx == 0? "NG":"OK"  std::to_string(b.det_score)
  //       cv::putText(img, cla_info.class_idx == 0? "NG":"OK", \
  //         cv::Point(b.det_box.x, b.det_box.y - 17), cv::FONT_HERSHEY_SIMPLEX, \
  //         1, 
  //         cla_info.class_idx == 0 ? cv::Scalar(0, 0, 255) : cv::Scalar(0, 255, 0) ,
  //         2);
  //     }
  //     if (count > 700 && count < 1200){
  //       writer << img;
  //     }
  //     if (count > 1200){
  //       break;
  //     }
  //     cv::imshow("img", img);
  //     cv::waitKey(1);
  //   }
  //   writer.release();
  // }

  // // text det && rec
  // {
  //   int ret = 0;
  //   // det
  //   std::string img_path(argv[2]);
  //   cv::Mat cv_image = cv::imread(img_path); 
  //   std::vector < Textai::DetInfo > text_lines;
  //   ret = pdmodel_manager->text_detect(&cv_image, text_lines, 0);
  //   std::cout << "text_detect ret: " << ret << std::endl;
  //   std::cout << "text_lines.size(): " << text_lines.size() << std::endl;
  //   // rec
  //   // std::vector<ai::TextRecInfo > ocr_infos;
  //   for (auto det_info : text_lines){
  //     ai::TextRecInfo crnn_info;
  //     ret = pdmodel_manager->text_recognise(&cv_image, det_info, 
  //                                         crnn_info, 1);
  //     if (ret != 0){continue;}
  //     std::cout << crnn_info.text << std::endl;
  //     // ocr_infos.push_back(crnn_info);
  //   }
  // }



  //   auto end = std::chrono::system_clock::now();
  //   auto duration =
  //       std::chrono::duration_cast<std::chrono::microseconds>(end - start);
  //   std::cout << "花费了"
  //             << double(duration.count()) *
  //                   std::chrono::microseconds::period::num /
  //                   std::chrono::microseconds::period::den
  //             << "秒" << std::endl;
  // }

  return 0;
}
